import argparse
import sys
import os
import gzip
import mmap
import numpy as np
import time
from multiprocessing import Manager
import multiprocessing as mp
from deepsignal3.utils.process_utils import get_motif_seqs
from deepsignal3.utils.process_utils import get_refloc_of_methysite_in_motif

neg_num = mp.Value('i', 0)

def read_feature(input,feature_Q,batch_size=1,qsize_limit=40):
    print('begin read feature', flush=True)
    fea_list=[]
    with open(input, 'r') as f:
        for line in f:
            words=line.strip().split('\t')
            key='||'.join(words[:2])
            while feature_Q.qsize()>=qsize_limit:
                time.sleep(0.01)
            fea_list.append(line)         
            if len(fea_list)>=batch_size:
                feature_Q.put(fea_list)
                fea_list=[]
    if len(fea_list)>0:
        feature_Q.put(fea_list)

    
def process_read(feature_Q,output_Q,standard,batch_size=1,time_wait=0.01,output_limit=40):
    fea_list=[]
    while True:
        while feature_Q.empty():
            time.sleep(time_wait)
        read_data = feature_Q.get()
        if read_data == "kill":
            feature_Q.put("kill")
            break
        for read in read_data:
            while output_Q.qsize()>=output_limit:
                time.sleep(0.01)
            words=read.strip().split('\t')

            #key='||'.join(words[:2])
            bisulfite=[float(x) for x in words[15].split(",")]
            density_results=words[-5:]
            density=[list(map(float, row.split(','))) for row in density_results]
            smaller_than_10 = [index for index, value in enumerate(bisulfite) if value <= 10]
            change=0
            for index in smaller_than_10:
                flag=0
                for s in range(len(standard)):
                    if density[s][index]!=-1 and density[s][index] <= standard[s]:
                        flag+=1
                if flag>=3:
                    bisulfite[index]=0
                    change+=1
            if change==0:
                continue
            bisulfite= [-1 if x != 0 else x for x in bisulfite]
            neg=[index for index, value in enumerate(bisulfite) if value == 0]
            neg_num.value+=len(neg)
            bisulfite_text = ','.join(
                        [str(bisulfite[x]) for x in range(len(bisulfite))])            
            fea_str='\t'.join(words[:15]+[bisulfite_text]+[words[16]])
            fea_list.append(fea_str)

            if len(fea_list)>=batch_size:
                output_Q.put(fea_list)
                fea_list=[]
    if len(fea_list)>0:
        output_Q.put(fea_list)

def _write_featurestr(write_fp, featurestr_q,time_wait = 3):
    print('write_process-{} starts'.format(os.getpid()))
    output=0
    # if os.path.exists(write_fp):
    #     with open(write_fp, 'a') as wf:
    #         while True:
    #             # during test, it's ok without the sleep(time_wait)
    #             if featurestr_q.empty():
    #                 time.sleep(time_wait)
    #                 continue
    #             features_str = featurestr_q.get()
    #             if features_str == "kill":
    #                 print('output line {}'.format(output))
    #                 print('write_process-{} finished'.format(os.getpid()))
    #                 break
    #             for one_features_str in features_str:
    #                 output+=1
    #                 wf.write(one_features_str + "\n")
    #             wf.flush()
    # else:
    with open(write_fp, 'w') as wf:
        while True:
            # during test, it's ok without the sleep(time_wait)
            if featurestr_q.empty():
                time.sleep(time_wait)
                continue
            features_str = featurestr_q.get()
            if features_str == "kill":
                print('output line {}'.format(output))
                print('write_process-{} finished'.format(os.getpid()))
                break
            for one_features_str in features_str:
                output+=1
                wf.write(one_features_str + "\n")
            wf.flush()

def parse_args():
    parser = argparse.ArgumentParser("")
    parser.add_argument("--input", type=str,required=True)
    parser.add_argument("--output", type=str,required=True)
    parser.add_argument("--nproc", "-p", type=int,required=True)
    return parser.parse_args()


def main():
    args=parse_args()
    #flag='w'
    #if os.path.exists(args.output):
    #    flag='a'

    manager = mp.Manager()
    feature_Q = manager.Queue()
    output_Q = manager.Queue()
    motif_seqs = get_motif_seqs('CG',True)
    ranges = list(range(10, 51, 10))
    standard=[0.0477,0.0489,0.0329,0.0248,0.0298]
  
    pd=mp.Process(target=read_feature,args=(args.input,feature_Q),name="fea_reader")
    pd.daemon = True
    pd.start()
    ex_dp=args.nproc - 2
    ex_procs = []
    for i in range(ex_dp):
        pb = mp.Process(target=process_read, args=(feature_Q,output_Q,standard),
                          name="pb_reader")
            
        pb.daemon = True
        pb.start()
        ex_procs.append(pb)
    p_w = mp.Process(target=_write_featurestr, args=(args.output, output_Q),
                     name="writer")
    p_w.daemon = True
    p_w.start()
    pd.join()
    feature_Q.put("kill")
    for pb in ex_procs:
        pb.join()
    output_Q.put("kill")
    p_w.join()
    print('add neg number is: {}'.format(neg_num.value),flush=True)
    
        


if __name__ == '__main__':
    sys.exit(main())

